/*
 *Author: Panos Sakkos
 *Email: p.sakkos@di.uoa.gr
 */

#include <mpi.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#define MASTER_RANK 0

typedef enum boolean {false = 0, true} bool;

void ReadMatrices(char *, int ***, int ***, int *);
void PrintMatrix(char *, int **, int);
int WorkersNeeded(int, int);
void InitMatrix(int ***, int);
int **GetBlock(int **, int, int, int);
void SetBlock(int ***, int **, int, int, int);
int **MultiplyMatrices(int **, int **, int);
int *PacketBlock(int **, int);
int **UnpacketBlock(int *, int);
bool Collector(int, int, int);
int GetSize(char *);
int MyTeamId(int , int , int );

int main(int argc, char **argv)
{
	int processes, myRank;	
	int size, numberOfTeams;
        int blockSize = atoi(argv[2]);

	if(argc != 3)
        {
    		fprintf(stderr, "Please, give only two arguments which indicate the input file and the number of blocks!\n");

                return EXIT_FAILURE;
        }

	size = GetSize(argv[1]);
	numberOfTeams = (size / blockSize) * (size * blockSize);

	MPI_Init(NULL, NULL);
	MPI_Comm_size(MPI_COMM_WORLD, &processes);
	MPI_Comm_rank(MPI_COMM_WORLD, &myRank);

	/* Create Intracommunicators for each team */

	MPI_Comm teamIntracommunicator;

	if(myRank != 0 )
	{
		MPI_Comm_split(MPI_COMM_WORLD, MyTeamId(myRank, size, blockSize), 0 , &teamIntracommunicator);
	}
	else
	{
		MPI_Comm_split(MPI_COMM_WORLD, MPI_UNDEFINED, 0 , &teamIntracommunicator);
	}

	if(myRank == 0)
	{
		/* Master's code */

	        int **matrixA, **matrixB, **matrixC;

	        ReadMatrices(argv[1], &matrixA, &matrixB, &size);
		InitMatrix(&matrixC, size);		

		/* Check if with the given matrix size and block size, the matrix 
		 * can be partitioned.
		 */
		
		int blocksCount = (size * size) / (blockSize * blockSize);
		
		if(size % blockSize != 0)
		{
			fprintf(stderr, "Cannot partition the %dx%d matrix  in %d %dx%d blocks!\n", 
						size, size, size / blockSize, blockSize, blockSize);

			MPI_Finalize();
			return EXIT_FAILURE;
		}

		/* Check if there are enough procceses spawned to do the work */

		else if( WorkersNeeded(size, blockSize)  != processes - 1)
		{
			fprintf(stderr, "Cannot run with a master and  %d workers for matrix size %d and block size %d.\n",
					processes - 1, size, blockSize);
			fprintf(stderr, "Master and %d workers needed!\n", WorkersNeeded(size, blockSize));

			MPI_Finalize();
			return EXIT_FAILURE;
		}
		else
		{
			printf("Partitioning in %d blocks matrix A and B\n", blocksCount);

			/* Create and share partitions */
		
			int i, j, k, workerId;
			
			workerId = 1;
			for(i = 0; i < size; i += blockSize)
			{
				for(j = 0; j < size; j += blockSize)
				{
					/* C[i,j] block is the result of i-th row of blocks of matrix A
					 * and j-th column of blocks of matrix B.
					 */
		
					for(k = 0; k < size; k += blockSize )
					{
						int **blockA = GetBlock(matrixA, i, k, blockSize);
						int **blockB = GetBlock(matrixB, k, j, blockSize);

						int *packetA = PacketBlock(blockA, blockSize);
						int *packetB = PacketBlock(blockB, blockSize);

						MPI_Send(packetA, blockSize * blockSize, MPI_INT, workerId, MASTER_RANK, MPI_COMM_WORLD);
						MPI_Send(packetB, blockSize * blockSize, MPI_INT, workerId, MASTER_RANK, MPI_COMM_WORLD);
						workerId++;
					}
				}
			}

			printf("Receiving blocks from every team's collectors\n");

			MPI_Status status;
			int receiveRank = 1;
			int *packetResult = malloc(sizeof(int) * blockSize * blockSize);

			/* Collect blocks from every team's collector */

			for(i = 0; i < size; i += blockSize)
			{
				for(j = 0; j < size; j += blockSize)
				{
					MPI_Recv(packetResult, blockSize * blockSize, MPI_INT, receiveRank, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
					int **resultBlock = UnpacketBlock(packetResult, blockSize);
					SetBlock(&matrixC, resultBlock, i, j, blockSize);
					receiveRank += size / blockSize;
				}
			}

			PrintMatrix("result.data", matrixC, size);
			printf("Result matrix computed. See context in result.data\n");
		}
	}
	else
	{
		/* Slave's code */

		MPI_Status status;

		int *packetA = malloc(sizeof(int) * blockSize * blockSize);
		MPI_Recv(packetA, blockSize * blockSize, MPI_INT, 0, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
		int **blockA = UnpacketBlock(packetA, blockSize);

		int *packetB = malloc(sizeof(int) * blockSize * blockSize);
                MPI_Recv(packetB, blockSize * blockSize, MPI_INT, 0, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
		int **blockB = UnpacketBlock(packetB, blockSize);

		int ** blockC = MultiplyMatrices(blockA, blockB, blockSize);
		int *packetC = PacketBlock(blockC, blockSize);
                
		/* Packet and reduce blockC with sum operation */

		int *packetResult = malloc(sizeof(int) * blockSize * blockSize);
		MPI_Reduce(packetC, packetResult, blockSize * blockSize, MPI_INT, MPI_SUM, 0, teamIntracommunicator);

		if(Collector(myRank, size, blockSize) == true)
		{
			MPI_Send(packetResult, blockSize * blockSize, MPI_INT, MASTER_RANK, myRank, MPI_COMM_WORLD);
		}
	}

	MPI_Finalize();
	return EXIT_SUCCESS;
}

int GetSize(char *inputFile)
{
        FILE *input = fopen(inputFile, "r");
                                                
        if(input == NULL)
        {
                perror("opening input file");
                return;
        }
                 
        /* Read dimension */
                                                
        int size;
        fscanf(input, "%d", &size);
	return size;
}

void ReadMatrices(char *inputFile, int ***matrixA, int ***matrixB, int *outSize)
{
	FILE *input = fopen(inputFile, "r");

	if(input == NULL)
	{
		perror("opening input file");
		return;
	}

	/* Read dimension */

	int size;
	fscanf(input, "%d", &size);
	*outSize = size;

	/* Allocate memory for matrices */

	*matrixA = malloc(sizeof(int *) * size);
	*matrixB = malloc(sizeof(int *) * size);

	int i, j;
	for(i = 0; i < size; i++ )
	{
		(*matrixA)[i] = malloc(sizeof(int) * size);
		(*matrixB)[i] = malloc(sizeof(int) * size);
	} 

	for(i = 0; i < size; i++)
	{
		for(j = 0; j < size; j++)
		{
			fscanf(input, "%d", &(*matrixA)[i][j]);
		} 
	}

	for(i = 0; i < size; i++)
	{
		for(j = 0; j < size; j++)
		{
			fscanf(input, "%d", &(*matrixB)[i][j]);
		}
	}
	
	fclose(input);
}

void PrintMatrix(char *filename, int **matrix, int size)
{
	FILE *file = fopen(filename, "w");

	int i, j;

	for(i = 0; i < size; i++)
	{
		for(j = 0; j < size; j++ )
		{
			fprintf(file, "%3d ", matrix[i][j]);
		}
		fprintf(file, "\n");
	}
	
	fprintf(file, "\n");

	fclose(file);
}

int WorkersNeeded(int size, int blockSize)
{
	return (size / blockSize ) * (size / blockSize) * (size / blockSize);
}

void InitMatrix(int ***matrix, int size)
{
	int i, j;

	*matrix = malloc(sizeof(int *) * size);

	for(i = 0; i < size; i++)
	{
		(*matrix)[i] = malloc(sizeof(int) * size);		

		for(j = 0; j < size; j++)
		{
			(*matrix)[i][j] = 0;
		}
	}
}

int **GetBlock(int **matrix, int iBlockStart, int jBlockStart, int size)
{
	int **block;

	block = malloc(sizeof(int *) * size);

	int i, j, blocki, blockj;
	for(i = iBlockStart, blocki = 0; i < iBlockStart + size; i++, blocki++)
	{
                block[blocki] = malloc(sizeof(int) * size);
		for(j = jBlockStart, blockj = 0; j < jBlockStart + size; j++, blockj++)
		{
			block[blocki][blockj] = matrix[i][j];
		}
	}

	return block;
}

void SetBlock(int ***matrix, int **block, int iBlockStart, int jBlockStart, int size)
{
	int i, j, blocki, blockj;
	for(i = iBlockStart, blocki = 0; blocki < size; i++, blocki++)
	{
		for(j = jBlockStart, blockj = 0; blockj < size; j++, blockj++)
		{
			(*matrix)[i][j] = block[blocki][blockj];
		}
	}
}

int **MultiplyMatrices(int **matrixA, int **matrixB, int size)
{
	int **result = malloc(sizeof(int *) * size);

	int i, j, k;
	for(i = 0; i < size; i++)
	{
		result[i] = malloc(sizeof(int) * size);
		for(j = 0; j < size; j++)
		{
			int sum = 0;
			for(k = 0; k < size; k++)
			{
				sum += matrixA[i][k] * matrixB[k][j];
			}
	
			result[i][j] = sum;
		}
	}

	return result;
}

int *PacketBlock(int **block, int size)
{
	int *packet = malloc(sizeof(int) * size * size);

	int i, j;	
	for(i = 0; i < size; i++)
	{
		for(j = 0; j < size; j++)
		{
			packet[i * size + j] = block[i][j];
		}
	}

	return packet;
}

int **UnpacketBlock(int *packet, int blockSize)
{
	int **block = malloc(sizeof(int *) * blockSize);

	int i, j;
	for(i = 0; i < blockSize; i++)
	{
		block[i] = malloc(sizeof(int) * blockSize);
		for(j = 0; j < blockSize; j++)
		{
			block[i][j] = packet[i * blockSize + j];
		}
	}

	return block;
}

bool Collector(int rank, int size, int blockSize)
{	
	return (rank - 1) % (size / blockSize) == 0 ;
}

int MyTeamId(int myRank, int size, int blockSize)
{
	return (myRank - 1) / (size / blockSize);
}

